#include "clang/AST/ASTConsumer.h"
#include "clang/AST/RecursiveASTVisitor.h"
#include "clang/Frontend/CompilerInstance.h"
#include "clang/Frontend/FrontendAction.h"
#include "clang/Tooling/Tooling.h"
#include "clang/Tooling/CommonOptionsParser.h"
#include "clang/Basic/DiagnosticOptions.h"
#include "clang/Frontend/TextDiagnosticBuffer.h"
#include "clang/Driver/Job.h"
#include "clang/Driver/Compilation.h"
#include "clang/Driver/Driver.h"
#include "clang/Driver/Tool.h"
#include "clang/Basic/TargetInfo.h"
#include "clang/Basic/TargetOptions.h"
#include "clang/../../lib/CodeGen/CGRecordLayout.h"
#include "clang/../../lib/CodeGen/CodeGenModule.h"
#include <llvm/ADT/ArrayRef.h>
#include <llvm/ADT/SmallString.h>
#include <llvm/Support/raw_ostream.h>

#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
#include "mlir/Dialect/GPU/IR/GPUDialect.h"
#include "mlir/IR/MLIRContext.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/BuiltinOps.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Pass/PassRegistry.h"
#include "mlir/Pass/PassManager.h"
#include "mlir/Tools/mlir-opt/MlirOptMain.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
#include "mlir/Transforms/DialectConversion.h"
#include "mlir/Transforms/Passes.h"
#include "mlir/InitAllDialects.h"
#include "mlir/InitAllPasses.h"

#include "crypto/Crypto.h"
#include "sst/SST.h"
#include "hyper/Hyper.h"

#include "pragmaHandler.hpp"

using namespace clang;
using namespace mlir;

namespace mlirclang {
/// Class to pass options to a compilation tool.
class ArgumentList {
private:
  /// Helper storage.
  llvm::SmallVector<llvm::SmallString<0>> Storage;
  /// List of arguments
  llvm::SmallVector<const char *> Args;

public:
  /// Add argument.
  ///
  /// The element stored will not be owned by this.
  void push_back(const char *Arg) { Args.push_back(Arg); }

  /// Add argument and ensure it will be valid before this passer's destruction.
  ///
  /// The element stored will be owned by this.
  template <typename... ArgTy> void emplace_back(ArgTy &&...Args) {
    // Store as a string
    std::string Buffer;
    llvm::raw_string_ostream Stream(Buffer);
    (Stream << ... << Args);
    emplace_back(llvm::StringRef(Stream.str()));
  }

  void emplace_back(llvm::StringRef &&Arg) {
    push_back(Storage.emplace_back(Arg).c_str());
  }

  /// Return the underling argument list.
  ///
  /// The return value of this operation could be invalidated by subsequent
  /// calls to push_back() or emplace_back().
  llvm::ArrayRef<const char *> getArguments() const { return Args; }
};
} // end namespace mlirclang

std::string GetExecutablePath(const char *Argv0, bool CanonicalPrefixes) {
  if (!CanonicalPrefixes) {
    SmallString<128> ExecutablePath(Argv0);
    // Do a PATH lookup if Argv0 isn't a valid path.
    if (!llvm::sys::fs::exists(ExecutablePath))
      if (llvm::ErrorOr<std::string> P =
              llvm::sys::findProgramByName(ExecutablePath))
        ExecutablePath = *P;
    return std::string(ExecutablePath.str());
  }

  // This just needs to be some symbol in the binary; C++ doesn't
  // allow taking the address of ::main however.
  void *P = (void *)(intptr_t)GetExecutablePath;
  return llvm::sys::fs::getMainExecutable(Argv0, P);
}

bool isNear(unsigned LineStart){
  for(size_t i = 0; i < wafer_guard::pragmaLineNumber.size(); ++i)
  {
    if(LineStart == wafer_guard::pragmaLineNumber[i] + 1) 
      return true;
  }
  return false;
}

class WaferConsumer : public clang::ASTConsumer 
{
public:

  explicit WaferConsumer(ASTContext *Context, Preprocessor &PP, mlir::OwningOpRef<mlir::ModuleOp> &module)
    : module(module)
    {
      addPragmaWaferHandlers(PP);
    }

  mlir::Type getMLIRType(clang::QualType qt, bool *implicitRef = nullptr, bool allowMerge = true);
  virtual void HandleTranslationUnit(clang::ASTContext &Context);

private:
  mlir::OwningOpRef<mlir::ModuleOp> &module;
};

class WaferVisitor : public RecursiveASTVisitor<WaferVisitor> 
{
public:
  explicit WaferVisitor(ASTContext *Context, WaferConsumer &Glob, mlir::OwningOpRef<mlir::ModuleOp> &module)
    : Context(Context), module(module), builder(module->getContext()), Glob(Glob) {}

  bool VisitCallExpr(clang::CallExpr *expr);

private:
  mlir::OwningOpRef<mlir::ModuleOp> &module;
  mlir::OpBuilder builder;
  ASTContext *Context;
  WaferConsumer &Glob;
};

class WaferAction : public clang::ASTFrontendAction 
{
public:

  mlir::OwningOpRef<mlir::ModuleOp> &module;

  WaferAction(mlir::OwningOpRef<mlir::ModuleOp> &module) : module(module) {}

  virtual std::unique_ptr<clang::ASTConsumer> CreateASTConsumer(clang::CompilerInstance &Compiler, llvm::StringRef InFile) {
    return std::unique_ptr<clang::ASTConsumer>(new WaferConsumer(&Compiler.getASTContext(), 
                                                                  Compiler.getPreprocessor(), module));
  }
};